package com.suyeer.basic.util;

import com.suyeer.basic.bean.BaseHttpResContent;
import com.suyeer.basic.bean.RSA;
import org.apache.commons.codec.binary.Base64;

import javax.crypto.Cipher;
import javax.servlet.http.HttpServletRequest;
import java.security.KeyFactory;
import java.security.PrivateKey;
import java.security.PublicKey;
import java.security.spec.PKCS8EncodedKeySpec;
import java.security.spec.X509EncodedKeySpec;

import static com.suyeer.basic.util.ConstUtil.KEY_ALGORITHM;
import static com.suyeer.basic.util.ConstUtil.PUBLIC_KEY;
import static com.suyeer.basic.util.ConstUtil.SESSION_RSA_KEY_IN_BASIC;

/**
 * @author jun 2018/10/19
 */
public class RSAUtil {

    /**
     * 公钥加密
     *
     * @param data         String
     * @param publicKeyStr String
     * @return String
     * @throws Exception Exception
     */
    public static String encrypt(String data, String publicKeyStr) throws Exception {
        byte[] bytes = Base64.decodeBase64(publicKeyStr);
        KeyFactory factory = KeyFactory.getInstance(KEY_ALGORITHM);
        X509EncodedKeySpec x509EncodedKeySpec = new X509EncodedKeySpec(bytes);
        PublicKey publicKey = factory.generatePublic(x509EncodedKeySpec);
        return encrypt(data, publicKey);
    }

    /**
     * 公钥加密
     *
     * @param data      String
     * @param publicKey PublicKey
     * @return String
     * @throws Exception Exception
     */
    public static String encrypt(String data, PublicKey publicKey) throws Exception {
        byte[] tmp1 = data.getBytes();
        Cipher cipher = Cipher.getInstance(KEY_ALGORITHM);
        cipher.init(Cipher.ENCRYPT_MODE, publicKey);
        byte[] tmp2 = cipher.doFinal(tmp1);
        return Base64.encodeBase64String(tmp2);
    }

    /**
     * 利用私钥解密
     *
     * @param base64Str     String
     * @param privateKeyStr String
     * @return String
     * @throws Exception Exception
     */
    public static String decrypt(String base64Str, String privateKeyStr) throws Exception {
        byte[] bytes = Base64.decodeBase64(privateKeyStr);
        KeyFactory factory = KeyFactory.getInstance(KEY_ALGORITHM);
        PKCS8EncodedKeySpec pkcs8EncodedKeySpec = new PKCS8EncodedKeySpec(bytes);
        PrivateKey privateKey = factory.generatePrivate(pkcs8EncodedKeySpec);
        return decrypt(base64Str, privateKey);
    }

    /**
     * 利用私钥解密
     *
     * @param base64Str  String
     * @param privateKey PrivateKey
     * @return String
     * @throws Exception Exception
     */
    public static String decrypt(String base64Str, PrivateKey privateKey) throws Exception {
        byte[] tmp1 = Base64.decodeBase64(base64Str);
        Cipher cipher = Cipher.getInstance(KEY_ALGORITHM);
        cipher.init(Cipher.DECRYPT_MODE, privateKey);
        byte[] tmp2 = cipher.doFinal(tmp1);
        return new String(tmp2);
    }

    /**
     * 获取 RSA 公钥并保存私钥
     *
     * @param request HttpServletRequest
     * @param attributeName String
     */
    public static void saveRsaInSession(HttpServletRequest request, String attributeName) {
        try {
            RSA rsa = SessionUtil.get(request, SESSION_RSA_KEY_IN_BASIC);
            if (rsa == null) {
                rsa = new RSA();
                SessionUtil.set(request, SESSION_RSA_KEY_IN_BASIC, rsa);
            }
            request.setAttribute(attributeName, rsa.getPublicKey());
        } catch (Exception e) {
            LogUtil.error(e.getMessage());
        }
    }

    /**
     * 获取 RSA 公钥并保存私钥
     *
     * @param request HttpServletRequest
     */
    public static void saveRsaInSession(HttpServletRequest request) {
        saveRsaInSession(request, PUBLIC_KEY);
    }

    /**
     * 解密, 私钥从session中获取
     *
     * @param base64Str String
     * @param hrc       BaseHttpResContent
     * @return String
     * @throws Exception Exception
     */
    public static String decryptWithPrivateKeyInSession(String base64Str, BaseHttpResContent hrc) throws Exception {
        return decryptWithPrivateKeyInSession(base64Str, hrc.getRequest());
    }

    /**
     * 解密, 私钥从session中获取
     *
     * @param base64Str String
     * @param request   HttpServletRequest
     * @return String
     * @throws Exception Exception
     */
    public static String decryptWithPrivateKeyInSession(String base64Str, HttpServletRequest request) throws Exception {
        RSA rsa = SessionUtil.get(request, SESSION_RSA_KEY_IN_BASIC);
        if (rsa == null) {
            throw new Exception("从 Session 中提取 RSA 秘钥对失败!");
        }
        return RSAUtil.decrypt(base64Str, rsa.getPrivateKey());
    }

    public static void main(String[] arg) {
        try {
            String data = "This is test data!";
            RSA rsaKey = new RSA();
            //公钥加密
            String tmp1 = encrypt(data, rsaKey.getPublicKey());
            System.err.println("\n加密后：\n\n" + tmp1);

            //私钥解密
            String tmp2 = decrypt(tmp1, rsaKey.getPrivateKey());
            System.err.println("\n解密后：\n\n" + tmp2);
        } catch (Exception e) {
            LogUtil.error(e);
        }
    }

}
